import torch.nn as nn


class DisplaceInterpolator(nn.Module):
    def __init__(self, input_size=1, hidden_size=128, num_layers=1):
        super().__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers=num_layers, batch_first=True)
        self.linear = nn.Linear(hidden_size, 40)
    
    def forward(self, modalinfo, hidden):
        out, hidden = self.lstm(modalinfo, hidden)  # out shape: (batch, seq_len, hidden_size)
        preds = self.linear(out[:, -1, :])
        
        return preds, hidden  # 只取最后一个时间步